package worker

import (
	"context"
	"github.com/sirupsen/logrus"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"simPSapp/pb"
)

type Worker struct {
	workerId         int
	serverTargetAddr string
	dataSize         int
	batchSize        int
	gradientSize     int // 单位MB

	id int
}

func NewWorker(serverAddr string, dataSize, batchSize, gradientSize, id int) *Worker {
	return &Worker{
		serverTargetAddr: serverAddr,
		dataSize:         dataSize,
		batchSize:        batchSize,
		gradientSize:     gradientSize,
		workerId:         id,
	}
}

func (w *Worker) Run() {
	conn, err := grpc.NewClient(w.serverTargetAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		logrus.Fatalf("[Run]-无法连接到服务器: %v", err)
	}
	defer conn.Close()
	client := pb.NewSendDataClient(conn)

	_, err = client.WorkerSync(context.Background(), &pb.HelloRequest{From: int32(w.workerId)})
	if err != nil {
		logrus.Fatalf("[Run]-无法同步: %v", err)
	}

	// 在数据集划分均匀的情况下，上传批次一般会相同。
	// 服务端在同步训练模式下，会等待所有客户端同步一轮梯度后再进行下一轮训练。

	payloadSize := w.dataSize*1024*1024/w.batchSize + 1 // 防止去尾后数据不够，这里+1

	for i := 0; i < w.batchSize; i++ {
		payload := make([]byte, payloadSize)

		// sendNByte回阻塞，直到其他worker都上传了本批次的数据
		_, err := client.SendNByte(context.Background(),
			&pb.RequestBytes{
				Payload: payload,
				From:    int32(w.workerId),
				BatchId: int32(i),
			})

		if err != nil {
			logrus.Fatalf("[Run]-无法上传数据: %v", err)
		}
	}

}
